import mitmproxy.http
from mitmproxy import ctx


class ListenWebSocket:
    def websocket_handshake(self, flow: mitmproxy.http.HTTPFlow):
        '''
        调用时间：客户端试图建立一个 websocket 连接。
        可以通过控制 HTTP 头部中针对 websocket 的条目来改变握手行为。flow 的 request 属性保证是非空的的。
        :param flow:
        :return:
        '''
        pass

    def websocket_start(self, flow: mitmproxy.websocket.WebSocketFlow):
        '''
        调用时间：建立了一个 websocket 连接。
        :param flow:
        :return:
        '''
        pass

    def websocket_message(self, flow: mitmproxy.websocket.WebSocketFlow):
        '''
        调用时间：收到一条来自客户端或服务端的 websocket 消息。最近一条消息存于 flow.messages[-1]。
        消息是可修改的。
        目前有两种消息类型，对应 BINARY 类型的 frame 或 TEXT 类型的 frame。
        :param flow:
        :return:
        '''
        pass

    def websocket_error(self, flow: mitmproxy.websocket.WebSocketFlow):
        '''
        调用时间：发生了 websocket 错误。
        :param flow:
        :return:
        '''
        pass

    def websocket_end(self, flow: mitmproxy.websocket.WebSocketFlow):
        '''
        调用时间：websocket 连接关闭。
        :param flow:
        :return:
        '''
        pass


addons = [
    ListenWebSocket()
]
